package com.sikaryofficial.gateway.filter;

import com.sikaryofficial.common.core.utils.ServletUtils;
import com.sikaryofficial.gateway.config.RateLimitProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.text.MessageFormat;
import java.time.Duration;

/**
 * @author : qinjinyuan
 * @desc : TODO  请填写你的功能描述
 * @date : 2024/08/14 15:26
 */
@Component
@Slf4j
public class RateLimitingFilterFactory extends AbstractGatewayFilterFactory<Object> {

    @Resource
    private RateLimitProperties rateLimitProperties;

    private static final String GATEWAY_RATE_LIMIT_KEY = "hismk:gateway:rate:limit:{0}";

    @Resource
    private ReactiveRedisTemplate<Object, Object> reactiveRedisTemplate;

    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            ServerHttpRequest request = exchange.getRequest();
            String path = request.getURI().getPath(); // 获取请求路径
            path = path.replace("//", "_");
            // 生成键
            String cacheKey = MessageFormat.format(GATEWAY_RATE_LIMIT_KEY, path);
            // 获取当前计数
            return reactiveRedisTemplate.opsForValue().increment(cacheKey, 1L)
                    .flatMap(count -> {
                        // 如果计数超过最大值，则拒绝请求
                        if (count > rateLimitProperties.getMaxRequestsPerSecond()) {
                            return handleTooManyRequests(exchange);
                        } else {
                            // 设置过期时间
                            return reactiveRedisTemplate.expire(cacheKey, Duration.ofSeconds(rateLimitProperties.getExpirationTime()))
                                    .then(chain.filter(exchange));
                        }
                    });
        };
    }

    private Mono<Void> handleTooManyRequests(ServerWebExchange exchange) {
        // 返回429 Too Many Requests HTTP状态码
        log.error("[Too Many Requests]请求路径:{}", exchange.getRequest().getPath());
        return ServletUtils.webFluxResponseWriter(exchange.getResponse(), HttpStatus.TOO_MANY_REQUESTS, "Too Many Requests",
                HttpStatus.TOO_MANY_REQUESTS.value());

    }
}